package com.googlecode.connectlet.machine.engine;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import com.googlecode.connectlet.Connection;
import com.googlecode.connectlet.Connector;
import com.googlecode.connectlet.ServerConnection;
import com.googlecode.connectlet.crypto.cipher.mode.CbcPkcs5PaddingDecipheror;
import com.googlecode.connectlet.crypto.cipher.mode.CbcPkcs5PaddingEncipheror;
import com.googlecode.connectlet.crypto.cipher.mode.EcbNoPaddingDecipheror;
import com.googlecode.connectlet.crypto.cipher.mode.EcbNoPaddingEncipheror;
import com.googlecode.connectlet.crypto.cipher.rc5.RC564Cipher;
import com.googlecode.connectlet.crypto.random.RC4PrngCipher;
import com.googlecode.connectlet.crypto.rsa.RSA;
import com.googlecode.connectlet.crypto.rsa.provider.CompactKeySpec;
import com.googlecode.connectlet.crypto.rsa.provider.RSAPrivateCrtKeyImpl;
import com.googlecode.connectlet.machine.MachinePacket;
import com.googlecode.connectlet.packet.PacketFilter;
import com.googlecode.connectlet.util.Bytes;
import com.googlecode.connectlet.util.ClassPath;
import com.googlecode.connectlet.util.db.ConnectionPool;
import com.googlecode.connectlet.util.db.DBExecutor;
import com.googlecode.connectlet.util.db.Row;

class MachineConnection extends Connection {
	class ResponseEvent extends Event {
		byte[] response;

		ResponseEvent(byte[] response) {
			super(MachinePacket.HEAD_TAG);
			this.response = response;
		}
	}

	private ExecutorService executor;

	MachineConnection(ExecutorService executor) {
		this.executor = executor;
		appendFilter(new PacketFilter(MachinePacket.getParser()));
	}

	@Override
	protected void onRecv(byte[] b, int off, int len) {
		final MachinePacket packet = new MachinePacket(b, off);
		if (packet.getRequestTag() == MachinePacket.TAG_PING) {
			send(new MachinePacket(MachinePacket.TAG_OK).getEncoded());
		} else {
			executor.execute(new Runnable() {
				@Override
				public void run() {
					byte[] response = MachineEngine.process(packet).getEncoded();
					dispatchEvent(new ResponseEvent(response), true);
				}
			});
		}
	}

	@Override
	protected void onEvent(Event event) {
		if (event instanceof ResponseEvent) {
			send(((ResponseEvent) event).response);
		}
	}
}

class MachineServer extends ServerConnection {
	private ExecutorService executor = Executors.newCachedThreadPool();

	MachineServer(int port) throws Exception {
		super(port);
	}

	@Override
	protected Connection createConnection() {
		return new MachineConnection(executor);
	}

	@Override
	public void close() {
		super.close();
		executor.shutdown();
	}
}

public class MachineEngine {
	static volatile boolean running;

	private static DBExecutor db;

	public static void startServer() throws Exception {
		Class.forName("org.h2.Driver");
		ClassPath dbFile = ClassPath.getInstance("../data/machine");
		final ConnectionPool connectionPool = new ConnectionPool("jdbc:h2:" + dbFile);
		db = new DBExecutor(connectionPool);
		if (db.query(1, "SELECT key_id FROM machine_key WHERE key_id = 0").error()) {
			db.source(ClassPath.getInstance("../data/MACHINE.SQL"));
		}
		final Connector connector = new Connector();
		connector.add(new MachineServer(8341));
		running = true;
		new Thread() {
			@Override
			public void run() {
				while (running) {
					while (connector.doEvents()) {/**/}
					try {
						sleep(10);
					} catch (InterruptedException e) {/**/}
				}
				connector.close();
				connectionPool.clear();
			}
		}.start();
	}

	public static void stopServer() {
		running = false;
	}

	private static RC4PrngCipher random = RC4PrngCipher.getInstance();

	private static int insertKey(byte[] key) {
		int[] id = new int[1];
		if (db.update(id, "INSERT INTO machine_key (key_data) VALUES (?)", key) != 1) {
			return -1;
		}
		return id[0];
	}

	private static void deleteKey(int id) {
		db.update("DELETE FROM machine_key WHERE key_id = ?", id);
	}

	private static byte[] getKey(int id) {
		Row row = db.query(1, "SELECT key_data FROM machine_key WHERE key_id = ?", id);
		return (row.valid() ? row.getBytes(1) : null);
	}

	private static final MachinePacket RESPONSE_ERROR = new MachinePacket(MachinePacket.TAG_ERROR);

	private static MachinePacket responseKeyId(int keyId) {
		return new MachinePacket(MachinePacket.TAG_OK, keyId);
	}

	private static MachinePacket responseKey(byte[] key) {
		return new MachinePacket(MachinePacket.TAG_OK, 0, 0, key);
	}

	private static MachinePacket responseData(byte[] data) {
		return new MachinePacket(MachinePacket.TAG_OK, 0, 0, Bytes.EMPTY_BYTES, data);
	}

	private static MachinePacket responseKeyPair(byte[] publicKey, byte[] privateKey) {
		return new MachinePacket(MachinePacket.TAG_OK, 0, 0, publicKey, privateKey);
	}

	static MachinePacket process(MachinePacket request) {
		switch (request.getRequestTag()) {
		case MachinePacket.TAG_KEY_NEW:
			return responseKeyId(insertKey(random.output(32)));
		case MachinePacket.TAG_KEY_SET:
			return responseKeyId(insertKey(request.getKey()));
		case MachinePacket.TAG_KEY_DEL:
			deleteKey(request.getKeyId());
			return new MachinePacket(MachinePacket.TAG_OK);
		}

		byte[] key = getKey(request.getKeyId());
		if (key == null) {
			return RESPONSE_ERROR;
		}

		switch (request.getRequestTag()) {
		case MachinePacket.TAG_KEY_WRAP:
		case MachinePacket.TAG_KEYPAIR_NEW:
			EcbNoPaddingEncipheror enc = new EcbNoPaddingEncipheror(new RC564Cipher(key));
			if (request.getRequestTag() == MachinePacket.TAG_KEY_WRAP) {
				return responseKey(enc.doFinal(request.getKey()));
			}
			if (request.getKeySize() < 512) {
				return RESPONSE_ERROR;
			}
			RSAPrivateCrtKeyImpl privateKey = new RSAPrivateCrtKeyImpl(request.getKeySize());
			return responseKeyPair(privateKey.getPublic().getEncoded(),
					enc.doFinal(privateKey.getEncoded()));
		case MachinePacket.TAG_PRIVATE_RSADP:
		case MachinePacket.TAG_SECRET_ENC_CBC:
		case MachinePacket.TAG_SECRET_DEC_CBC:
			EcbNoPaddingDecipheror dec = new EcbNoPaddingDecipheror(new RC564Cipher(key));
			if (request.getRequestTag() == MachinePacket.TAG_PRIVATE_RSADP) {
				byte[] encodedKey = dec.doFinal(request.getKey());
				privateKey = new CompactKeySpec(encodedKey).getPrivateKey();
				return responseData(RSA.rsadp(privateKey, request.getData()));
			}
			byte[] data = request.getData();
			if (data.length < 16) {
				return RESPONSE_ERROR;
			}
			byte[] key2 = dec.doFinal(request.getKey());
			if (request.getRequestTag() == MachinePacket.TAG_SECRET_ENC_CBC) {
				byte[] iv = random.output(16);
				CbcPkcs5PaddingEncipheror cbc = new CbcPkcs5PaddingEncipheror(new RC564Cipher(key2), iv);
				return responseData(Bytes.add(iv, cbc.doFinal(data)));
			}
			CbcPkcs5PaddingDecipheror cbc = new CbcPkcs5PaddingDecipheror(new RC564Cipher(key2), data);
			return responseData(cbc.doFinal(data, 16, data.length - 16));
		}

		return RESPONSE_ERROR;
	}
}